/**
 * Copyright 2009 Erick Armbrust 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0 
 *   
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 */

package tapioca.dt;

import java.util.List;
import java.util.Map;

import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.repackaged.com.google.common.base.Preconditions;
import com.google.appengine.repackaged.com.google.common.collect.Lists;

/**
 * A {@link DatastoreService} wrapper that provides functionality for reading, writing, and
 * deleting DT enabled objects.
 *
 * @author armbrust@gmail.com (Erick Armbrust)
 */
public class WrappedDatastoreService {

  public static final String WRITE_LOCK_PROPERTY = "dt__write_lock";
  public static final String VERSION_PROPERTY = "dt__version";
  public static final String NAMED_KEY_PREFIX = "dt__";  
  public static final String PMD_KIND = "DT__PureMetaData__";
  public static final String PMD_KEY_NAME = "pmd"; 
  
  private final DatastoreService ds;
  
  /**
   * Constructor.
   * 
   * @param ds the {@link DatastoreService} to wrap.
   */
  public WrappedDatastoreService(DatastoreService ds) {
    this.ds = ds;
  }  
  
  /**
   * Retrieves an entity, populates an {@link GetResponse}, checks to see if the entity is DT
   * capable, and returns.  This method returns null for the entity if the entity is not found.
   * 
   * @param key the key of the item to retrieve from the datastore.
   * @return the {@link GetResponse} for the given key.
   */
  public GetResponse get(Key key) {
    return get(null, key);
  }  
  
  /**
   * Retrieves an entity in a transaction, populates an {@link GetResponse}, checks to see if the 
   * entity is DT capable, and returns.  This method returns null for the entity if the entity 
   * is not found.
   * 
   * @param txn the transaction to use for this request, or null if a transaction should be created.
   * @param key the key of the item to retrieve from the datastore.
   * @return the {@link GetResponse} for the given key.
   */  
  public GetResponse get(Transaction txn, Key key) {
    // Only commit transactions that we start locally.
    boolean commitPendingTransaction = (txn == null);
    
    // Create a transaction if necessary.
    txn = (txn == null) ? ds.beginTransaction() : txn;
    
    Entity entity = null;
    Entity pureMetaData = null;
    Key writeLock = null;
    Key version = null;
   
    // Standard case for non-named keys.  Simply fetch the object from the datastore and
    // extract the write lock and version.
    if (key.getName() == null) {
      try {
        entity = ds.get(txn, key);
        writeLock = (Key) entity.getProperty(WRITE_LOCK_PROPERTY);
        version = (Key) entity.getProperty(VERSION_PROPERTY);
      } catch (EntityNotFoundException e) {
        // This is okay, the caller can figure out if this entity exists by checking the
        // the GetResponse.
      }
      
    // In the case of a named key, the object must be fetched in a transaction along with its
    // metadata entity.
    } else {
      Key pureMetaDataKey = createPureMetaDataKey(key);
      Map<Key, Entity> entityAndPmd = ds.get(txn, Lists.newArrayList(key, pureMetaDataKey));      
      
      // If an entity is found in the datastore, its metadata object must exist as well.
      if (entityAndPmd.containsKey(key) && !entityAndPmd.containsKey(pureMetaDataKey)) {
        txn.rollback();
        throw new IllegalStateException("Named entity exists but its metadata does not");
      }
      
      // Get the entity if it exists.
      if (entityAndPmd.containsKey(key)) {
        entity = entityAndPmd.get(key);
      }
      
      // Get the metadata if it exists, or create it if it does not exist.
      if (entityAndPmd.containsKey(pureMetaDataKey)) {
        pureMetaData = entityAndPmd.get(pureMetaDataKey);
        writeLock = (Key) pureMetaData.getProperty(WRITE_LOCK_PROPERTY);
        version = (Key) pureMetaData.getProperty(VERSION_PROPERTY);
      } else {
        pureMetaData = createPureMetaData(key);
      }
    }
    
    // Only commit the transaction if it was started within this method.
    if (commitPendingTransaction && txn.isActive()) {
      txn.commit();
    }    
    
    // Before returning, make sure this is distributed transaction capable.
    checkIfDtCapable(entity, pureMetaData);
    return new GetResponse(entity, pureMetaData, writeLock, version);    
  }
  
  /**
   * Puts an item into the datastore.  Since all DT entities must have a version, the version is
   * updated within this method.  When an entity is put to the datastore, it's write lock is
   * cleared.
   * 
   * @param txn the transaction to use for this put operation.
   * @param entity the entity to be persisted to the datastore.
   * @param version the new version for the entity.
   */
  public void put(Transaction txn, Entity entity, Key version) {
    Entity pureMetaData = entity.getKey().getName() == null ? null : 
        createPureMetaData(entity.getKey());
    Entity toUpdate = pureMetaData == null ? entity : pureMetaData;
    toUpdate.setProperty(WRITE_LOCK_PROPERTY, null);
    toUpdate.setProperty(VERSION_PROPERTY, version);
    
    List<Entity> toPut = Lists.newArrayList(entity);
    if (pureMetaData != null) toPut.add(pureMetaData);
    ds.put(txn, toPut);    
  }
  
  public void put(Transaction txn, List<Entity> entities, Key version) {
    List<Entity> toPut = Lists.newArrayList(entities);
    
    for (Entity entity : entities) {
      Entity pureMetaData = entity.getKey().getName() == null ? null : 
          createPureMetaData(entity.getKey());
      Entity toUpdate = pureMetaData == null ? entity : pureMetaData;
      toUpdate.setProperty(WRITE_LOCK_PROPERTY, null);
      toUpdate.setProperty(VERSION_PROPERTY, version);
      
      if (pureMetaData != null) { 
        toPut.add(pureMetaData);
      }
      
      toPut.add(entity);
    }
    
    if (toPut.size() > 0) {
      ds.put(txn, toPut);
    }
  }
  
  /**
   * Deletes an item from the datastore.  Named entities must have a version persisted at all times,
   * even if the object is deleted.
   * 
   * @param txn the transaction to use for the delete operation.
   * @param key the key for the item that is being deleted.
   * @param version the new version for the entity.
   */
  public void delete(Transaction txn, Key key, Key version) {
    Entity pureMetaData = key.getName() == null ? null : createPureMetaData(key);    
    
    if (pureMetaData != null) {
      pureMetaData.setProperty(WRITE_LOCK_PROPERTY, null);
      pureMetaData.setProperty(VERSION_PROPERTY, version);
      ds.put(txn, pureMetaData);
    }
    
    ds.delete(txn, key);
  }
  
  public void delete(Transaction txn, List<Key> keys, Key version) {
    List<Entity> toPut = Lists.newArrayList();
    
    for (Key key : keys) {
      Entity pureMetaData = key.getName() == null ? null : createPureMetaData(key);    
      
      if (pureMetaData != null) {
        pureMetaData.setProperty(WRITE_LOCK_PROPERTY, null);
        pureMetaData.setProperty(VERSION_PROPERTY, version);
        toPut.add(pureMetaData);
      }      
    }
    
    if (toPut.size() > 0) {
      ds.put(txn, toPut);
    }
    
    if (keys.size() > 0) {
      ds.delete(txn, keys);
    }
  }
  
  /**
   * Updates the write lock for a given entity.
   * 
   * @param txn the transaction to use for the write lock put operation.
   * @param response the {@link GetResponse} that will provide the entity as well as metadata.
   * @param writeLock the new write lock for the entity.
   */
  public void setAndSaveWriteLock(Transaction txn, GetResponse response, Key writeLock) {
    Entity toSave = response.getPureMetaData() == null ? response.getEntity() : 
        response.getPureMetaData();
    if (toSave != null) {
      toSave.setProperty(WRITE_LOCK_PROPERTY, writeLock);
      ds.put(txn, toSave);
    }
  }  
  
  ///// Static functions /////
  
  /**
   * Creates a key for a pure metadata entity used to store write lock/version information for
   * entities with named keys.
   * 
   * @param key the named key for which we should create a pure metadata key. 
   * @return the generated pure metadata key.
   */
  public static Key createPureMetaDataKey(Key key) {
    Preconditions.checkArgument(key.isComplete() && key.getName() != null, 
        "Cannot create a pure metadata key for a non-named key");
    return KeyFactory.createKey(key, PMD_KIND, PMD_KEY_NAME);
  }
  
  /**
   * Creates a full pure metadata entity given a named key.
   * 
   * @param key the named key for which we should create a pure metadata entity.
   * @return the pure metadata entity.
   */
  public static Entity createPureMetaData(Key key) {
    Entity pureMetaData = new Entity(createPureMetaDataKey(key));
    pureMetaData.setProperty(WRITE_LOCK_PROPERTY, null);
    pureMetaData.setProperty(VERSION_PROPERTY, null);
    return pureMetaData;
  }
  
  /**
   * Checks to see if an entity is distributed transaction capable.  What this means is that the
   * entity must have the following properties:
   * 
   * Non-named entities:
   * For non-named entities, which are entities without a named key, the entity must have both 
   * write lock and version properties.
   * 
   * Named entities:
   * For named entities, which are entities with a named key, the entity must have a pure metadata
   * object which contains the write lock and version properites. 
   * 
   * @param entity the entity to inspect.
   * @param pureMetaData the pure metadata entity for that object if applicable.
   */
  public static void checkIfDtCapable(Entity entity, Entity pureMetaData) {
    Entity toCheck = entity;
    if (pureMetaData != null) {
      toCheck = pureMetaData;
    }
    
    if (toCheck != null) {
      Preconditions.checkArgument(toCheck.hasProperty(WRITE_LOCK_PROPERTY),
          "Entity does not have the " + WRITE_LOCK_PROPERTY + " property, " +
          "not a valid DT entity.");
      Preconditions.checkArgument(toCheck.hasProperty(VERSION_PROPERTY),
          "Entity does not have the " + VERSION_PROPERTY + " property, " +
          "not a valid DT entity.");
    } 
    
    if (entity != null) {
      Key key = entity.getKey();
      if (key.getName() != null) {
        Preconditions.checkArgument(key.getName().startsWith(NAMED_KEY_PREFIX),
            "Entity with named key " + key.getName() + " does not have the " +
            "correct prefix \"" + NAMED_KEY_PREFIX + "\"");    
      }
    }
  }  
  
  /**
   * A helper class that represents the results of a DT entity get.
   */
  public static class GetResponse {
    
    private final Entity entity;
    private final Entity pureMetaData;
    private final Key writeLock;
    private final Key version;
    
    /**
     * A private constructor since the {@link WrappedDatastoreService}'s {@code get} method should
     * be the only way to instantiate a get response.
     */
    private GetResponse(Entity entity, Entity pureMetaData, Key writeLock, Key version) {
      this.entity = entity;
      this.pureMetaData = pureMetaData;
      this.writeLock = writeLock;
      this.version = version;
    }

    public boolean inDatastore() {
      return entity != null;
    }
    
    public Entity getEntity() {
      return entity;
    }

    public Entity getPureMetaData() {
      return pureMetaData;
    }

    public Key getWriteLock() {
      return writeLock;
    }

    public Key getVersion() {
      return version;
    }
  }  
  
}
